package main

import (
	"context"
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"tunnel_manager/core"

	"github.com/google/uuid"
	"github.com/quic-go/quic-go"
)

// ClientPooledConnection 客户端池化的QUIC连接
type ClientPooledConnection struct {
	Conn          *quic.Conn   // 🔧 使用*quic.Conn类型
	UDPConn       *net.UDPConn // 🔥 底层UDP连接（需要手动管理生命周期）
	ConnectionID  string       // 连接唯一ID
	ProtocolType  ProtocolType // 协议类型
	ActiveStreams int32        // 当前活跃Stream数（原子操作）
	TotalStreams  int64        // 历史总Stream数（原子操作）
	LastUsedTime  time.Time    // 最后使用时间
	mu            sync.RWMutex // 保护LastUsedTime
}

// IncrementStreams 增加活跃Stream计数
func (c *ClientPooledConnection) IncrementStreams() {
	atomic.AddInt32(&c.ActiveStreams, 1)
	atomic.AddInt64(&c.TotalStreams, 1)
	c.mu.Lock()
	c.LastUsedTime = time.Now()
	c.mu.Unlock()
}

// DecrementStreams 减少活跃Stream计数
func (c *ClientPooledConnection) DecrementStreams() {
	atomic.AddInt32(&c.ActiveStreams, -1)
}

// GetActiveStreams 获取当前活跃Stream数
func (c *ClientPooledConnection) GetActiveStreams() int32 {
	return atomic.LoadInt32(&c.ActiveStreams)
}

// GetTotalStreams 获取历史总Stream数
func (c *ClientPooledConnection) GetTotalStreams() int64 {
	return atomic.LoadInt64(&c.TotalStreams)
}

// ClientQUICConnectionPool 客户端QUIC连接池
type ClientQUICConnectionPool struct {
	serverAddr   string
	tunnelID     string
	tlsConfig    *tls.Config
	quicConfig   *quic.Config
	poolSize     int                                             // 连接池大小
	connections  map[core.ProtocolType][]*ClientPooledConnection // 按协议类型分组
	mu           sync.RWMutex
	ctx          context.Context
	cancel       context.CancelFunc
	reconnecting map[string]bool // 正在重连的连接ID
	reconnectMu  sync.Mutex
	roundRobin   map[core.ProtocolType]*int32 // 🚀 轮询索引（每个协议一个）
}

// NewClientQUICConnectionPool 创建客户端QUIC连接池
func NewClientQUICConnectionPool(serverAddr, tunnelID string, tlsConfig *tls.Config, quicConfig *quic.Config, poolSize int) *ClientQUICConnectionPool {
	ctx, cancel := context.WithCancel(context.Background())
	pool := &ClientQUICConnectionPool{
		serverAddr:   serverAddr,
		tunnelID:     tunnelID,
		tlsConfig:    tlsConfig,
		quicConfig:   quicConfig,
		poolSize:     poolSize,
		connections:  make(map[core.ProtocolType][]*ClientPooledConnection),
		ctx:          ctx,
		cancel:       cancel,
		reconnecting: make(map[string]bool),
		roundRobin:   make(map[core.ProtocolType]*int32), // 🚀 初始化轮询索引
	}

	// 🔥🔥🔥 暂时禁用健康检查
	// 问题：Context().Done()不是可靠的连接状态指示器
	// QUIC库内部的Context可能因为各种原因被取消，但连接实际上还是活跃的
	// 解决方案：移除主动健康检查，只在实际使用时发现连接不可用才重连
	// go pool.healthCheck() // ❌ 禁用

	return pool
}

// ConnectAll 建立所有连接
func (p *ClientQUICConnectionPool) ConnectAll(protocolType core.ProtocolType) error {
	log.Printf("🔗 [连接池] 开始建立连接: 协议=%s, 数量=%d", protocolType, p.poolSize)

	var wg sync.WaitGroup
	errors := make(chan error, p.poolSize)

	for i := 0; i < p.poolSize; i++ {
		wg.Add(1)
		go func(index int) {
			defer wg.Done()
			if err := p.addNewConnection(protocolType, index); err != nil {
				errors <- err
			}
		}(i)
	}

	wg.Wait()
	close(errors)

	// 检查错误
	var firstError error
	errorCount := 0
	for err := range errors {
		if firstError == nil {
			firstError = err
		}
		errorCount++
	}

	if errorCount > 0 {
		log.Printf("⚠️  [连接池] 部分连接失败: %d/%d, 错误: %v", errorCount, p.poolSize, firstError)
		// 只要有至少1个连接成功就继续
		if errorCount == p.poolSize {
			return fmt.Errorf("所有连接失败: %v", firstError)
		}
	}

	successCount := p.poolSize - errorCount
	log.Printf("✅ [连接池] 连接建立完成: 成功=%d, 失败=%d", successCount, errorCount)
	return nil
}

// addNewConnection 添加新连接
func (p *ClientQUICConnectionPool) addNewConnection(protocolType ProtocolType, index int) error {
	connID := uuid.New().String()

	log.Printf("🔌 [连接池] 建立连接#%d: ConnID=%s, 协议=%s", index, connID[:8]+"...", protocolType)

	// 建立QUIC连接
	udpAddr, err := net.ResolveUDPAddr("udp", p.serverAddr)
	if err != nil {
		return fmt.Errorf("解析服务器地址失败: %v", err)
	}

	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		return fmt.Errorf("创建UDP连接失败: %v", err)
	}
	// 🔥🔥🔥 关键修复：不要立即关闭UDP连接！
	// UDP连接是QUIC的底层传输，需要在QUIC连接的整个生命周期内保持打开
	// 只有当QUIC连接关闭时，UDP连接才应该关闭
	// defer udpConn.Close() // ❌ 删除此行，否则QUIC连接会立即失效

	conn, err := quic.Dial(p.ctx, udpConn, udpAddr, p.tlsConfig, p.quicConfig)
	if err != nil {
		return fmt.Errorf("建立QUIC连接失败: %v", err)
	}

	// 发送认证消息
	stream, err := conn.OpenStreamSync(p.ctx)
	if err != nil {
		conn.CloseWithError(0, "认证失败")
		return fmt.Errorf("打开认证Stream失败: %v", err)
	}

	authMsg := &TunnelMessage{
		Magic:        ProtocolMagic,
		Version:      ProtocolVersion,
		MessageType:  MessageTypeRequest,
		ProtocolType: protocolType,
		TunnelID:     p.tunnelID,
		ConnectionID: connID,
		RequestID:    "", // 认证消息不需要RequestID
		BodyLength:   0,
		Body:         nil,
	}

	if err := authMsg.Write(stream); err != nil {
		stream.Close()
		conn.CloseWithError(0, "认证失败")
		return fmt.Errorf("发送认证消息失败: %v", err)
	}

	// 等待服务端响应
	respMsg, err := ReadMessage(stream)
	if err != nil {
		stream.Close()
		conn.CloseWithError(0, "认证失败")
		return fmt.Errorf("读取认证响应失败: %v", err)
	}

	stream.Close()

	if respMsg.MessageType == MessageTypeError {
		conn.CloseWithError(0, "认证失败")
		return fmt.Errorf("认证失败: %s", string(respMsg.Body))
	}

	// 添加到连接池
	pooledConn := &ClientPooledConnection{
		Conn:          conn,    // 🔧 修复：conn已经是*quic.Conn类型
		UDPConn:       udpConn, // 🔥 保存UDP连接引用，用于后续清理
		ConnectionID:  connID,
		ProtocolType:  protocolType,
		ActiveStreams: 0,
		TotalStreams:  0,
		LastUsedTime:  time.Now(),
	}

	p.mu.Lock()
	p.connections[protocolType] = append(p.connections[protocolType], pooledConn)
	p.mu.Unlock()

	log.Printf("✅ [连接池] 连接#%d建立成功: ConnID=%s", index, connID[:8]+"...")
	return nil
}

// GetConnection 获取连接（🚀🚀🚀 视频分片优化：真正的负载均衡）
func (p *ClientQUICConnectionPool) GetConnection(protocolType core.ProtocolType) (*ClientPooledConnection, error) {
	// 🚀 优化1：只在读取连接列表时持有锁，然后立即释放
	p.mu.RLock()
	conns, exists := p.connections[protocolType]
	if !exists || len(conns) == 0 {
		p.mu.RUnlock()
		return nil, fmt.Errorf("没有可用的%s协议连接", protocolType)
	}

	// 🚀 优化2：复制连接切片，避免长时间持有锁
	connsCopy := make([]*ClientPooledConnection, len(conns))
	copy(connsCopy, conns)
	p.mu.RUnlock()

	// 🔥🔥🔥 关键优化：智能负载均衡（专为视频Range分片并发优化）
	// 策略：
	// 1. 永远使用轮询（Round-Robin）确保负载均衡
	// 2. 不优先选择"空闲"连接，因为这会导致多个Range请求集中到同一个连接
	// 3. 让每个新请求都分散到不同的连接上，实现真正的并发

	// 初始化轮询索引（如果不存在）
	p.mu.RLock()
	if p.roundRobin[protocolType] == nil {
		p.mu.RUnlock()
		p.mu.Lock()
		if p.roundRobin[protocolType] == nil {
			idx := int32(0)
			p.roundRobin[protocolType] = &idx
		}
		p.mu.Unlock()
		p.mu.RLock()
	}
	p.mu.RUnlock()

	// 🚀 使用原子操作进行轮询，确保线程安全且不阻塞
	currentIdx := atomic.AddInt32(p.roundRobin[protocolType], 1) - 1
	selectedIdx := int(currentIdx) % len(connsCopy)
	bestConn := connsCopy[selectedIdx]

	log.Printf("🎯 [客户端连接池] 轮询选择连接: ConnID=%s, 索引=%d/%d, 活跃Stream=%d",
		bestConn.ConnectionID[:8]+"...", selectedIdx+1, len(connsCopy), bestConn.GetActiveStreams())

	return bestConn, nil
}

// RemoveConnection 移除连接
func (p *ClientQUICConnectionPool) RemoveConnection(connID string, protocolType ProtocolType) {
	p.mu.Lock()
	defer p.mu.Unlock()

	conns, exists := p.connections[protocolType]
	if !exists {
		return
	}

	for i, conn := range conns {
		if conn.ConnectionID == connID {
			p.connections[protocolType] = append(conns[:i], conns[i+1:]...)
			// 🔥 先关闭QUIC连接，再关闭底层UDP连接
			(*conn.Conn).CloseWithError(0, "移除连接")
			if conn.UDPConn != nil {
				conn.UDPConn.Close()
			}
			log.Printf("🔌 [连接池] 移除连接: ConnID=%s, 协议=%s, 剩余=%d",
				connID[:8]+"...", protocolType, len(p.connections[protocolType]))
			break
		}
	}
}

// GetSize 获取指定协议的连接数
func (p *ClientQUICConnectionPool) GetSize(protocolType ProtocolType) int {
	p.mu.RLock()
	defer p.mu.RUnlock()

	if conns, exists := p.connections[protocolType]; exists {
		return len(conns)
	}
	return 0
}

// CleanupDisconnectedConnections 清理已断开的连接
// 🔥 智能清理策略：检查Context().Done()来快速识别断开的连接
func (p *ClientQUICConnectionPool) CleanupDisconnectedConnections(protocolType ProtocolType) int {
	p.mu.Lock()
	defer p.mu.Unlock()

	conns, exists := p.connections[protocolType]
	if !exists {
		return 0
	}

	healthyConns := make([]*ClientPooledConnection, 0, len(conns))
	removedCount := 0

	for _, conn := range conns {
		// 检查连接是否已断开
		select {
		case <-(*conn.Conn).Context().Done():
			// 连接已断开
			log.Printf("🧹 [客户端清理] 检测到断开的连接: ConnID=%s",
				conn.ConnectionID[:8]+"...")

			// 关闭连接
			(*conn.Conn).CloseWithError(0, "cleanup disconnected")
			if conn.UDPConn != nil {
				conn.UDPConn.Close()
			}

			removedCount++
		default:
			// 连接仍然活跃
			healthyConns = append(healthyConns, conn)
		}
	}

	// 更新连接列表
	p.connections[protocolType] = healthyConns

	if removedCount > 0 {
		log.Printf("✅ [客户端清理] 完成: 移除%d个断开的连接，剩余%d个健康连接",
			removedCount, len(healthyConns))
	}

	return removedCount
}

// GetStats 获取连接池统计信息
func (p *ClientQUICConnectionPool) GetStats(protocolType ProtocolType) ClientPoolStats {
	p.mu.RLock()
	defer p.mu.RUnlock()

	conns, exists := p.connections[protocolType]
	if !exists || len(conns) == 0 {
		return ClientPoolStats{}
	}

	var totalActive int32
	var totalHistorical int64
	var minActive, maxActive int32 = 999999, 0

	for _, conn := range conns {
		active := conn.GetActiveStreams()
		totalActive += active
		totalHistorical += conn.GetTotalStreams()

		if active < minActive {
			minActive = active
		}
		if active > maxActive {
			maxActive = active
		}
	}

	return ClientPoolStats{
		TotalConnections:   len(conns),
		TotalActiveStreams: totalActive,
		TotalStreams:       totalHistorical,
		MinActiveStreams:   minActive,
		MaxActiveStreams:   maxActive,
		AvgActiveStreams:   float64(totalActive) / float64(len(conns)),
	}
}

// healthCheck 健康检查，自动重连失败的连接
func (p *ClientQUICConnectionPool) healthCheck() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-p.ctx.Done():
			return
		case <-ticker.C:
			p.checkAndReconnect()
		}
	}
}

// checkAndReconnect 检查并重连失败的连接
func (p *ClientQUICConnectionPool) checkAndReconnect() {
	p.mu.Lock()
	defer p.mu.Unlock()

	for protocolType, conns := range p.connections {
		var validConns []*ClientPooledConnection
		needReconnect := 0

		for _, conn := range conns {
			// 🔥🔥🔥 修复：使用更可靠的方式检查连接状态
			// 问题：Context().Done()可能被过早触发
			// 解决方案：尝试打开一个Stream来测试连接是否真的活跃

			// 先检查Context（快速检查）
			select {
			case <-(*conn.Conn).Context().Done():
				// Context已取消，连接可能已关闭
				// 但再做一次确认：尝试打开Stream
				testCtx, testCancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
				stream, err := (*conn.Conn).OpenStreamSync(testCtx)
				testCancel()

				if err != nil {
					// 确实无法打开Stream，连接已死
					needReconnect++
					log.Printf("⚠️  [连接池] 检测到关闭的连接: ConnID=%s, error=%v", conn.ConnectionID[:8]+"...", err)
					// 关闭UDP连接
					if conn.UDPConn != nil {
						conn.UDPConn.Close()
					}
					continue
				} else {
					// 可以打开Stream，连接还活着！Context被误判了
					stream.Close()
					log.Printf("✅ [连接池] 连接实际上还活着: ConnID=%s (Context虽然Done但Stream正常)", conn.ConnectionID[:8]+"...")
					validConns = append(validConns, conn)
				}
			default:
				// Context正常，连接应该是好的
				validConns = append(validConns, conn)
			}
		}

		p.connections[protocolType] = validConns

		if needReconnect > 0 {
			log.Printf("🔄 [连接池] 需要重连: 协议=%s, 数量=%d", protocolType, needReconnect)
			// 异步重连
			go func(pt ProtocolType, count int) {
				for i := 0; i < count; i++ {
					if err := p.addNewConnection(pt, -1); err != nil {
						log.Printf("❌ [连接池] 重连失败: %v", err)
					}
				}
			}(protocolType, needReconnect)
		}
	}
}

// Close 关闭连接池
func (p *ClientQUICConnectionPool) Close() {
	p.cancel()

	p.mu.Lock()
	defer p.mu.Unlock()

	for _, conns := range p.connections {
		for _, conn := range conns {
			// 🔥 先关闭QUIC连接，再关闭底层UDP连接
			(*conn.Conn).CloseWithError(0, "连接池关闭")
			if conn.UDPConn != nil {
				conn.UDPConn.Close()
			}
		}
	}

	p.connections = make(map[ProtocolType][]*ClientPooledConnection)
	log.Printf("🛑 [连接池] 客户端连接池已关闭")
}

// ClientPoolStats 客户端连接池统计信息
type ClientPoolStats struct {
	TotalConnections   int     // 总连接数
	TotalActiveStreams int32   // 总活跃Stream数
	TotalStreams       int64   // 历史总Stream数
	MinActiveStreams   int32   // 最小活跃Stream数
	MaxActiveStreams   int32   // 最大活跃Stream数
	AvgActiveStreams   float64 // 平均活跃Stream数
}
